import sys
parent_dir = "../pygl"
sys.path.append(parent_dir)
import ctypes
import OpenGL.GL as gl
import glm
import imgui
import glfw
from glfw import _GLFWwindow as GLFWwindow
import numpy as np
import pygl
from common import AppWithTrackball, Sphere, AxesHelper, GridHelper

class LorenzAttractor(AppWithTrackball):
    # 初始化着色器，vao vbo等
    def setupVertex(self) -> None:
        gl.glEnable(gl.GL_DEPTH_TEST)
        gl.glEnable(gl.GL_BLEND)
        gl.glBlendFunc(gl.GL_SRC_ALPHA, gl.GL_ONE_MINUS_SRC_ALPHA)
        gl.glEnable(gl.GL_MULTISAMPLE)
        self.light_program = pygl.ProgramVF("intermediate/shaders/geometry.vs", 
                                      "intermediate/shaders/geometry.fs")
        self.program = pygl.ProgramVF("intermediate/shaders/lorenz.vs", 
                                      "intermediate/shaders/lorenz.fs")
        self.sphere_program = pygl.ProgramVF("intermediate/shaders/phong.vs", 
                                      "intermediate/shaders/phong.fs")
        self.background_color = glm.vec4(0.8, 0.8, 0.8, 1.0)
        self.grid = GridHelper(300.0, 30, glm.vec3(0.8, 0.8, 0.8), 
                               glm.vec3(0.2, 0.2, 0.2))
        self.show_grid = False
        self.world_axes = AxesHelper(100, 5.0)
        self.show_world_axes = False
        self.imgui_window_width = 250
        self.imgui_window_height = 200
        self.light = Sphere(0.5)
        self.light_color = glm.vec3(1.0, 1.0, 1.0)
        self.light.position = glm.vec3(0.0, 50.0, 50.0)
        self.sphere = Sphere(2.0)
        self.sphere_color = glm.vec3(46.0/255.0, 99.0/255.0, 190.0/255.0)
        self.sphere1 = Sphere(2.0)
        self.sphere1_color = glm.vec3(46.0/255.0, 99.0/255.0, 0.0)
        self.ambient_color = glm.vec3(1.0, 1.0, 1.0)
        self.Ka = 0.4
        self.Kd = 0.9
        self.Ks = 1.0
        self.shininess = 60.0
        # 设置变化量
        dt = 0.003
        # 设置计算步数
        self.num_steps = 50000
        # 保存每一步计算步数的x y z，定义为顶点数据
        self.vertices = np.empty((self.num_steps + 1, 6), dtype=gl.GLfloat)
        self.vertices1 = np.empty((self.num_steps + 1, 6), dtype=gl.GLfloat)
        # 定义两个起始点的位置，相差非常小
        self.vertices[0, 0] = 0.0
        self.vertices[0, 1] = 1.0
        self.vertices[0, 2] = 0.0
        self.vertices1[0, 0] = 0.0
        self.vertices1[0, 1] = 1.0
        self.vertices1[0, 2] = 0.00001
        for i in range(self.num_steps):
            dx, dy, dz = self.lorenz(self.vertices[i, 0], self.vertices[i, 1], self.vertices[i, 2])
            self.vertices[i + 1, 0] = self.vertices[i, 0] + dx * dt
            self.vertices[i + 1, 1] = self.vertices[i, 1] + dy * dt
            self.vertices[i + 1, 2] = self.vertices[i, 2] + dz * dt
            self.vertices[i + 1, 3] = self.vertices[i + 1, 0] / 5.0
            self.vertices[i + 1, 4] = self.vertices[i + 1, 1] / 5.0
            self.vertices[i + 1, 5] = self.vertices[i + 1, 2] / 25.0
            dx, dy, dz = self.lorenz(self.vertices1[i, 0], self.vertices1[i, 1], self.vertices1[i, 2])
            self.vertices1[i + 1, 0] = self.vertices1[i, 0] + dx * dt
            self.vertices1[i + 1, 1] = self.vertices1[i, 1] + dy * dt
            self.vertices1[i + 1, 2] = self.vertices1[i, 2] + dz * dt
            self.vertices1[i + 1, 3] = self.vertices1[i + 1, 0] / 5.0
            self.vertices1[i + 1, 4] = self.vertices1[i + 1, 1] / 5.0
            self.vertices1[i + 1, 5] = self.vertices1[i + 1, 2] / 25.0
        self.vao = pygl.VertexArrayObject()
        self.vbo = pygl.VertexBufferObject(self.vertices)
        attribute_position = pygl.VertexAttribute("position", 0, 3, gl.GL_FLOAT, gl.GL_FALSE, 0)
        attribute_color = pygl.VertexAttribute("color", 1, 3, gl.GL_FLOAT, gl.GL_FALSE, 3 * gl.sizeof(gl.GLfloat))
        self.vao.setVertexBuffer(self.vbo, 0, 0, 6 * gl.sizeof(gl.GLfloat))
        self.vao.setVertexAttribute(0, attribute_position)
        self.vao.setVertexAttribute(0, attribute_color)
        self.vao1 = pygl.VertexArrayObject()
        self.vbo1 = pygl.VertexBufferObject(self.vertices1)
        self.vao1.setVertexBuffer(self.vbo1, 0, 0, 6 * gl.sizeof(gl.GLfloat))
        self.vao1.setVertexAttribute(0, attribute_position)
        self.vao1.setVertexAttribute(0, attribute_color)
        self.model = glm.mat4(1.0)
        self.camera.far = 400
        self.camera.zoom_speed = 2.0
        self.camera.reset(glm.vec3(2.0, 1.0, 100.0), glm.vec3(0.0, 0.0, 0.0), 
                    glm.vec3(0.0, 1.0, 0.0))
        self.animation_step = 1

    # 配置imgui界面元素
    def showImGui(self)->None:
        imgui.begin("设置")
        if imgui.button("重置"):
            self.camera.reset(glm.vec3(2.0, 1.0, 100.0), glm.vec3(0.0, 0.0, 0.0), 
                    glm.vec3(0.0, 1.0, 0.0))
        _, self.show_world_axes = imgui.checkbox("显示三维世界坐标轴", 
                                                self.show_world_axes)
        _, self.show_grid = imgui.checkbox("显示XZ平面网格", 
                                                self.show_grid)
        imgui.end()       

    # 渲染图形
    def render(self) -> None:
        gl.glClearBufferfv(gl.GL_COLOR, 0, self.background_color.to_tuple())
        gl.glClearBufferfv(gl.GL_DEPTH, 0, 1.0)
        self.view = self.camera.getViewMatrix()
        self.projection = self.camera.getProjectionMatrix()
        # 代表灯位置的球体
        self.light_program.use()
        self.light.updateModelMatrix()
        self.light_program.setUniformMatrix4fv(0, self.light.model_matrix)
        self.light_program.setUniformMatrix4fv(1, self.view)
        self.light_program.setUniformMatrix4fv(2, self.projection)
        self.light_program.setUniform1i(3, 1)
        self.light_program.setUniform3fv(4, self.light_color)
        self.light.render()
        self.program.use()
        view = self.camera.getViewMatrix()
        projection = self.camera.getProjectionMatrix()
        self.program.setUniformMatrix4fv(0, self.model)
        self.program.setUniformMatrix4fv(1, view)
        self.program.setUniformMatrix4fv(2, projection)
        gl.glLineWidth(2.0)
        if self.animation_step <= self.num_steps:
            self.vao.bind()
            gl.glDrawArrays(gl.GL_LINE_STRIP, 0, self.animation_step)
            self.vao1.bind()
            gl.glDrawArrays(gl.GL_LINE_STRIP, 0, self.animation_step)
            self.animation_step += 1
        else:
            self.vao.bind()
            gl.glDrawArrays(gl.GL_LINE_STRIP, 0, self.num_steps)
            self.vao1.bind()
            gl.glDrawArrays(gl.GL_LINE_STRIP, 0, self.num_steps)
        self.sphere_program.use()
        self.vao.bind()
        self.sphere.position = glm.vec3(self.vertices[self.animation_step - 1, 0],
                            self.vertices[self.animation_step - 1, 1],
                            self.vertices[self.animation_step - 1, 2])
        self.sphere.updateModelMatrix()
        self.sphere_program.setUniformMatrix4fv(0, self.sphere.model_matrix)
        self.sphere_program.setUniformMatrix4fv(1, self.view)
        self.sphere_program.setUniformMatrix4fv(2, self.projection)
        self.sphere_program.setUniform3fv(3, self.sphere_color)
        self.sphere_program.setUniform3fv(4, self.light_color)
        self.sphere_program.setUniform3fv(5, self.light.position)
        self.sphere_program.setUniform3fv(6, self.ambient_color)
        self.sphere_program.setUniform3fv(7, self.camera.position)
        self.sphere_program.setUniform1f(8, self.Ka)
        self.sphere_program.setUniform1f(9, self.Kd)
        self.sphere_program.setUniform1f(10, self.Ks)
        self.sphere_program.setUniform1f(11, self.shininess)
        self.sphere_program.setUniform1i(12, 4)
        self.sphere.render()
        self.vao1.bind()
        self.sphere1.position = glm.vec3(self.vertices1[self.animation_step - 1, 0],
                            self.vertices1[self.animation_step - 1, 1],
                            self.vertices1[self.animation_step - 1, 2])
        self.sphere1.updateModelMatrix()
        self.sphere_program.setUniformMatrix4fv(0, self.sphere1.model_matrix)
        self.sphere_program.setUniformMatrix4fv(1, self.view)
        self.sphere_program.setUniformMatrix4fv(2, self.projection)
        self.sphere_program.setUniform3fv(3, self.sphere1_color)
        self.sphere_program.setUniform3fv(4, self.light_color)
        self.sphere_program.setUniform3fv(5, self.light.position)
        self.sphere_program.setUniform3fv(6, self.ambient_color)
        self.sphere_program.setUniform3fv(7, self.camera.position)
        self.sphere_program.setUniform1f(8, self.Ka)
        self.sphere_program.setUniform1f(9, self.Kd)
        self.sphere_program.setUniform1f(10, self.Ks)
        self.sphere_program.setUniform1f(11, self.shininess)
        self.sphere_program.setUniform1i(12, 4)
        self.sphere1.render()
        if self.show_world_axes:
            self.world_axes.updateModelMatrix
            self.world_axes.render(self.view, self.projection)
        if self.show_grid:
            self.grid.render(view, projection)

    # 释放资源
    def cleanup(self) -> None:
        self.light.delete()
        self.grid.delete()
        self.sphere.delete()
        self.sphere1.delete()
        self.world_axes.delete()
        self.light_program.delete()
        self.program.delete()
        self.sphere_program.delete()
        self.vao.delete()
        self.vao1.delete()
        self.vbo.delete()
        self.vbo1.delete()

    def mouse_move(self, window:GLFWwindow, x:float, y:float)->None:
        if imgui.is_any_item_active():
            return
        super().mouse_move(window, x, y)
    
    # 定义洛伦兹方程
    def lorenz(self, x : float, y : float, z : float, sigma = 10, rho = 28, beta = 2.667) -> tuple:
        dx = sigma * (y - x)
        dy = x * (rho - z) - y
        dz = x * y - beta * z
        return dx, dy, dz

if __name__ == "__main__":
    app = LorenzAttractor(100, 100, 1000, 600, "洛伦兹吸引子App！")
    app.run()


